"""用于在工具和推理步骤之间共享信息的上下文管理器。"""
import json
import time
from typing import Any, Dict, List, Optional
import uuid
from langgraph.store.memory import BaseStore

from model.context_model import ReasoningStep, ToolContext
from model.memory_model import MemoryType


class ContextManager:
    """管理工具与推理步骤之间的上下文和信息流。"""

    def __init__(self, memory_store: BaseStore):
        self.memory_store = memory_store
        self.current_session_id: Optional[str] = None
        self.reasoning_steps: List[ReasoningStep] = []
        self.tool_contexts: Dict[str, List[ToolContext]] = {}
        self.shared_variables: Dict[str, Any] = {}


    def start_session(self, session_id: str, initial_query: str):
        """开始一个新的推理会话。"""
        self.current_session_id = session_id
        self.reasoning_steps = []
        self.tool_contexts = {}

        self.shared_variables = {
            "initial_query": {
                "value": initial_query,
                "source_tool": None,
                "timestamp": time.time()
            }
        }

    async def add_reasoning_step(self, step: ReasoningStep):
        """向当前会话添加一个推理步骤。"""
        self.reasoning_steps.append(step)

        await self.memory_store.aput(
            (MemoryType.WORKING.value, self.current_session_id),
            key=str(step.step_number),
            value={
                'content': {
                    "step_number": step.step_number,
                    "thought": step.thought,
                    "planned_action": step.planned_action,
                    "confidence": step.confidence
                },
                'metadata': {
                    "session_id": self.current_session_id,
                    "step_type": "reasoning"
                },
                'importance': step.confidence
            }
        )

    async def add_tool_context(self, tool_context: ToolContext):
        """添加工具调用的上下文信息"""
        if tool_context.tool_name not in self.tool_contexts:
            self.tool_contexts[tool_context.tool_name] = []
        self.tool_contexts[tool_context.tool_name].append(tool_context)

        # 只保存工具调用成功的结果
        if tool_context.success:
            await self.memory_store.aput(
                (MemoryType.TOOL_CONTEXT.value, ),
                key=str(uuid.uuid4()),
                value={
                    "key": f'Action: {tool_context.tool_name}, Action Input: {tool_context.input_data}',
                    "tool_context": tool_context.model_dump_json(),
                },
                index=['key']
            )

    def set_shared_variable(self, key: str, value: Any, source_tool: Optional[str] = None):
        """设置一个可供其他工具使用的共享变量。"""
        self.shared_variables[key] = {
            "value": value,
            "source_tool": source_tool,
            "timestamp": time.time()
        }

    def get_shared_variable(self, key: str) -> Any:
        var_data = self.shared_variables.get(key)
        return var_data["value"] if var_data else None

    def get_all_shared_variables(self) -> Dict[str, Any]:
        return {k: v["value"] for k, v in self.shared_variables.items()}

    async def get_relevant_context(self, current_tool: str, query: str) -> Dict[str, Any]:
        """获取工具调用相关的上下文"""
        context = {
            "shared_variables": self.get_all_shared_variables(),
            "previous_tool_results": {},
            "reasoning_history": [],
            "similar_past_executions": []
        }

        # 获取上次工具调用的结果
        for tool_name, contexts in self.tool_contexts.items():
            if contexts:
                latest_context = contexts[-1]
                if latest_context.success:
                    context["previous_tool_results"][tool_name] = {
                        "output": latest_context.output_data,
                        "metadata": latest_context.metadata
                    }

        # 获取最近工具调用的结果
        context["reasoning_history"] = [
            {
                "thought": step.thought,
                "action": step.planned_action,
                "confidence": step.confidence
            }
            for step in self.reasoning_steps[-5:]
        ]

        # 检索过去类似的工具成功执行的记录
        similar_memories = await self.memory_store.asearch(
            (MemoryType.TOOL_CONTEXT.value, ),
            query=f"Action: {current_tool}, Action Input: {query}",
            limit=3
        )
        similar_memories = similar_memories or []
        similar_tool_contexts = [
            ToolContext(**json.loads(memory.value['tool_context'])) \
            for memory in similar_memories if memory is not None
        ]
        context["similar_past_executions"] = [
            {
                "tool_name": tool_context.tool_name,
                "input": tool_context.input_data,
                "output": tool_context.output_data,
                "success_rate": tool_context.importance,
            }
            for tool_context in similar_tool_contexts
        ]

        return context

    async def end_session(self):
        """结束当前会话并存储情景记忆。"""
        pass
        # if self.current_session_id:
        #     session_summary = {
        #         "session_id": self.current_session_id,
        #         "total_reasoning_steps": len(self.reasoning_steps),
        #         "tools_used": list(self.tool_contexts.keys()),
        #         "shared_variables_count": len(self.shared_variables),
        #     }
        #
        #     # 将整个会话存储为情景记忆
        #     await self.memory_store.aput(
        #         (MemoryType.EPISODIC.value, ),
        #         key=self.current_session_id,
        #         value={
        #             "session_summary": session_summary,
        #             "reasoning_steps": [
        #                 {
        #                     "thought": step.thought,
        #                     "action": step.planned_action,
        #                     "success": step.tool_context.success if step.tool_context else None
        #                 }
        #                 for step in self.reasoning_steps
        #             ],
        #             "final_shared_variables": self.get_all_shared_variables(),
        #             "importance": 0.7,
        #             "metadata": {
        #                 "session_id": self.current_session_id,
        #                 "tools_used": session_summary["tools_used"]
        #             }
        #         }
        #     )